

"""
2022.05.20 Liukb
    根据遗传日常变异核对需求，提取结果文件内的数据
    正常情况，运行速度在 5 秒内，依据 cdh 访问速度，会有卡顿和延迟。
"""

import os
import re
import sys
import requests
import pandas as pd
from var_info_collect import Var_Info
from collections import defaultdict
sys.path.insert(0, '/analysis_s140/liukb/my_script/python/SV')
from SV_vcf_tools import VCF


class SV_info:
    def __init__(self, sv_file: str) -> None:
        self.sv_file = sv_file
        self.sv = defaultdict(str)
        self.genes = defaultdict(list)
        self.parse_f()

    def parse_f(self):
        """ 解析 SV 注释文件 """
        if not os.path.isfile(self.sv_file):
            print(f'\n没有 SV 结果文件: {self.sv_file}')
            return
        with open(self.sv_file) as fi:
            for line in fi:
                cells = line.strip().split('\t')
                if line.startswith('confidence_score'):
                    confidence_score = cells.index('confidence_score')
                    Annotation_mode = cells.index('Annotation_mode')
                    SV_chrom = cells.index('SV_chrom')
                    SV_start = cells.index('SV_start')
                    SV_end = cells.index('SV_end')
                    SV_type = cells.index('SV_type')
                    SV_gene = cells.index('Gene_name')
                    ACMG_class = cells.index('ACMG_class')
                    continue
                # 分基因注释的跳过
                if cells[Annotation_mode] == 'split':
                    continue
                sv_info = f"{cells[SV_chrom]} {cells[SV_start]} {cells[SV_end]} {cells[SV_type]} {cells[confidence_score]} {cells[ACMG_class]}"
                self.sv[(cells[SV_chrom], cells[SV_start], cells[SV_end])] = sv_info
                for gene in cells[SV_gene].split(';'):
                    self.genes[gene].append(sv_info)

    def call_var(self, region: str):
        """ 查询未知相近的 SV """
        Chr, Pos = region.split(':')
        var_Start, var_End = Pos.split('-')
        res = ''
        for (Chrom, Start, End), var in self.sv.items():
            if Chr != Chrom:
                continue
            if VCF.cal_overlap(int(var_Start), int(var_End), int(Start), int(End)) >= 0.8:
                res = f"{res}{var}\n"
        return res

    def call_gene(self, gene_name: str):
        if gene_name in self.genes:
            return "\n".join(self.genes[gene_name])



class Search_var:
    def __init__(self):
        # self.sample_index_dict = self.build_index()
        self.fasta: str = '/analysis_s140/reference/GATK_bundle_b37/human_g1k_v37_decoy.fasta'

    # @staticmethod
    # def build_index():
    #     """
    #         构建样本编号与文件对应关系字典（提示 索引中有多个同名样本）
    #     """
    #     sample_dict = {}
    #     for sf in [
    #         '/cdh/ngs/analysis/huangsx/data_analysis_archive/archive_short_txt.txt',
    #         '/cdh/ngs/archive/archive_short_txt.txt'
    #     ]:
    #         with open(sf) as f:
    #             for line in f.readlines():
    #                 if any(s in line for s in ['/reanno/', '/old_anno/', '/new_anno/']):
    #                     continue
    #                 sample_ID = os.path.basename(line).split('.')[0]
    #                 batch = line.split('/analysis/report/')[0].split('/')[-1]
    #                 if sample_ID in sample_dict.keys():
    #                     d1 = sample_dict[sample_ID][0].split('/')[-3]
    #                     d2 = line.split('/')[-3]
    #                     if int(d2) > int(d1):
    #                         sample_dict[sample_ID][0] = '/cdh' + line.strip()
    #                     sample_dict[sample_ID][1].append(batch)
    #                 else:
    #                     sample_dict[sample_ID] = ['/cdh' + line.strip(), [batch]]
    #     return sample_dict
                        
    @staticmethod
    def clean_input(s):
        """
            清洗输入，确定输入有哪些类型的变异
            示例输入：
                WES22FJ0000874 JAK1 1-65345127 
                3 53529193 53529195 ，WES21FJ0003052，POLR3B，PI4KA
        """
        s = re.sub(r'@[^\s]+\s', ' ', s)
        s = re.sub(r'exoncnv', ' ', s, re.I)
        s = re.sub(r'\([^\)]+\)', ' ', s)
        s = re.sub(r'[\-_\W\u4e00-\u9fa5]+', ' ', s)
        s = s.strip()
        # 先查询位置性的变异
        sample_ID = ''
        chr_list = [str(i + 1) for i in range(22)] + ['X', 'Y']
        loc_list = []
        loc_i = -1
        loc = False
        gene_list = []
        match_index = 0
        for i, k in enumerate(s.split(' ')):
            if match_index < i:
                loc = False
            if loc:
                if k.isdigit() and k not in chr_list:
                    loc_list[loc_i] += f' {k}'
                    continue
            if k in chr_list:
                match_index = i + 2
                loc_i += 1
                loc_list.append(k)
                loc = True
                continue
            else:
                if len(k) == 14 and k[0:3] in ['NBS', 'SAG', 'DAT', 'WES', 'MES', 'EQA']:
                    sample_ID = k
                    continue
                else:
                    if len(k) > 1 and not k.isdigit():
                        gene_list.append(k)
                        continue
                    else:
                        print(f'========== 【警告】 这个不像基因，不知道是啥 {k}')
                        continue
        loc_list = list(set(loc_list))
        gene_list = [g for g in gene_list if not g.lower().startswith('exon')]
        gene_list = list(set(gene_list))
        return (sample_ID, loc_list, gene_list)

    @staticmethod
    def sample_info(s):
        url = 'http://10.10.10.150:8250/ngs/SamplePath'
        json_data = {
            'sample' : s,
            'new' : '0'
        }
        res = requests.post(url, json=json_data)
        if res.status_code != 200:
            print(f'{s} 没有在索引库内')
            return
        info = res.json()['message'][s]
        all_batch = ''
        local_df = {}
        if len(info) > 1:
            for batch, dt in info.items():
                local_df[dt['date']] = dt
                all_batch = f"{all_batch},{batch}"
            new_date = str(max([int(i) for i in local_df.keys()]))
            info = local_df[new_date]
            info['other_batch'] = all_batch.replace(info['batch_name'], '').strip(',')
        else:
            info = dict(*info.values())
            info['other_batch'] = ''
        return info

    def Run(self, s):
        """ 除了输入信息 """
        cnv_header = [
            'Chr', 'Start', 'End', 'SVSize', 'SVType', 'CytoBand', 'Flag', 
            'CopyNum', 'Probes', 'Gene', 'GeneCount', 'predict', 'pli_counts'
        ]
        sample_ID, loc_list, gene_list = self.clean_input(s)
        sample_index = self.sample_info(sample_ID)
        # print(sample_index)
        batch = sample_index['batch_name']
        self.Bam_file = sample_index['bam_file']
        exonCNV_file = sample_index['exon_cnv_txt']
        sv_file = exonCNV_file.replace(
            f'{sample_ID}.CNVExon.txt', 'Anno/AnnotSV.tsv'
        ).replace(
            'analysis/report/', 'analysis/variant/SV/'
        )
        self.sv_info = SV_info(sv_file)
        CNV_file = exonCNV_file.replace('CNVExon.txt', 'CNV.txt')
        stat_file = CNV_file.replace(f'{sample_ID}/{sample_ID}.CNV.txt', 'stat.txt')
        sex, name, r1 = self.close_depth_sample(stat_file, sample_ID)
        print(f'\n{"-"*50}\n\t样本编号: <{sample_ID}> {name} {sex}  {batch}\n\t位置: {loc_list}\n\t基因: {"，".join(gene_list)}')
        print(f'{"-"*50}\n{r1}\n{self.rel_info(sample_ID)}')
        res = ''
        # 点突变
        warns = ''
        for k in loc_list:
            try:
                Chr, Pos = k.split(' ', 1)
            except ValueError:
                warns = f'%%%% 无效信息: {k} %%%\n'
                continue
            if ' ' in Pos:
                Start, End = Pos.split(' ')
                # 大片段 CNV
                if int(End) - int(Start) > 200 and os.path.isfile(CNV_file):
                    res += self.CNV_search(CNV_file, cnv_header, Chr, Start, End)
                    continue
            res += self.loc_search(sample_index['snp_txt'], Chr, *Pos.split(' '))
        # exonCNV
        for gene in gene_list:
            res += self.exonCNV_search(exonCNV_file, gene)
            if os.path.isfile(CNV_file):
                res += self.CNV_search(CNV_file, cnv_header, gene_name=gene)
        print(res.strip(), end='\n\n')
        if sample_index['other_batch']:
            print(f'\n【注意】 样本存在于多个批次 : {sample_index["other_batch"]}\n{warns}')

    def loc_search(self, short_file, Chr, Start, End=''):
        """
            点突变查询，输出格式：
                1-9 列的基本信息
                HGVS:
                人群频率（仅列有结果的）：
                致病性预测（仅列有结果的）：
                致病性（仅列有结果的）：
        """
        def add_res(cells, key, end='\n'):
            try:
                value = cells[header.index(key)]
                if value != '.':
                    return f'[{key}]: {value} {end}'
                return ''
            except ValueError:
                print(f'{key} 不存在\t')
                return ''
        def split_s(s):
            o = ''
            if len(s) > 100:
                if re.search(r' \t ', s[100:]):
                    s_100 = s[0:100]
                    s1, s2 = s[100:].split(' \t ', 1)
                    o = f'{s_100}{s1}\n\t\t{split_s(s2)}'
                else:
                    o = s
            else:
                o = s
            return o
        if not End:
            End = Start
        out = ''
        # vars: list = []
        with open(short_file) as f:
            for line in f.readlines():
                if line.startswith('Chr'):
                    header = line.split('\t')
                    continue
                cells = line.split('\t')
                if Chr == cells[0] and Start == cells[1] and End == cells[2]:
                    var_info = Var_Info(
                        ' '.join(cells[0:5]),
                        self.Bam_file,
                        self.fasta
                    )
                    try:
                        var_info.run()
                    except:
                        var_info.score = ''
                        var_info.level = ''
                    int_val = ''
                    for k in ['source', 'fatherGT', 'motherGT', 'var_Confidence']:
                        if k in header:
                            int_val = f'{int_val}   {cells[header.index(k)]}'
                    res = f'=== \033[1;37;46m{"   ".join(cells[0:9])}   {int_val.strip()}\033[0m  {var_info.score}  \033[1;37;41m {var_info.level} \033[0m \n\t'
                    # res = f'=== \033[1;37;46m{"   ".join(cells[0:9])}   {int_val.strip()}\033[0m \n\t'
                    for k in ['GeneDetail.refGene', 'VEP_HGVSc', 'VEP_HGVSp']:
                        res += add_res(cells, k, '\n\t')
                    res = res.strip() + "\n\t人群频率: "
                    for k in [
                        'Han_Frequency', 'Inter_Freq_WES', 'X1000g2015aug_all', 'X1000g2015aug_eas', 
                        'ExAC_ALL', 'ExAC_EAS', 'AF_popmax', 'AF_sas', 'AF_eas', 'controls_AF_popmax', 
                        'gnomAD_genome_ALL', 'gnomAD_genome_EAS'
                    ]:
                        res += add_res(cells, k, ' \t ')
                    res = res.strip("\n\t人群频率: ").strip(' \t ') + "\n\t有害性预测: "
                    for k in [
                        'REVEL_score', 'ClinPred_Score', 'SIFT_pred', 'Polyphen2_HDIV_pred', 
                        'Polyphen2_HVAR_pred', 'LRT_pred', 'MutationTaster_pred', 
                        'MutationAssessor_pred', 'PROVEAN_pred', 'CADD_raw', 'CADD_phred', 
                        'dbscsnv11_AdaBoost', 'dbscsnv11_RandomForest', 'LoFtool', 
                        'SpliceAI_pred', 'MaxEntScan', 'GeneSplicer'
                    ]:
                        res += add_res(cells, k, ' \t ')
                    res = res.strip("\n\t有害性预测: ").strip(' \t ') + "\n\t评级: "
                    for k in ['CLNSIG', 'InterVar_automated']:
                        res += add_res(cells, k, ' \t ')
                    res = res.strip("\n\t评级: ").strip(' \t ') + "\n\t"
                    res += add_res(cells, 'pseudogene')
                    for s in re.findall(r'[^\n]+', res.strip()):
                        out = f'{out}\n{split_s(s)}'
        if out:
            return out.strip() + "\n"
        return f'\033[1;31;40m=== {Chr} {Start} {End} 变异未找到!\033[0m\n'

    @staticmethod
    def gene_loc(gene_name):
        """ 返回基因所有转录本范围并集的区间 """
        with open('/analysis_s140/reference/annovar_humandb/hg19_refGene.txt') as f1:
            Chr, Start, End = '', '0', '0'
            chrs = [str(i) for i in range(1,23)] + ['X', 'Y', 'MT']
            for line in f1.readlines():
                cells = line.split('\t')
                if gene_name == cells[12] and cells[2].replace('chr', '') in chrs:
                    if Start == '0' or int(cells[4]) < int(Start):
                        Start = cells[4]
                    if int(cells[5]) > int(End):
                        End = cells[5]
                    Chr = cells[2].replace('chr', '')
            return f"{Chr}  {Start}  {End}" if Chr else ''

    def exonCNV_search(self, exonCNV_file, gene_name):
        """
            exonCNV 查询：
                显示1-12 列信息
            未查询到时，从 /analysis_s140/reference/annovar_humandb/hg19_refGene.txt 内查询基因的位置（仅返回第一个即可）
        """
        SV_res = self.sv_info.call_gene(gene_name)
        with open(exonCNV_file, encoding='gbk') as f:
            for line in f.readlines():
                cells = line.split('\t')
                if gene_name == cells[6]:
                    cells[6] = f'\033[1;30;43m{cells[6]}\033[0m'
                    if SV_res:
                        info = f'>>> {"   ".join(cells[0:12])}\n\033[1;33;40mSV 检出\033[0m:{SV_res}\n' 
                    else:
                        info = f'>>> {"   ".join(cells[0:12])}\n\n' 
                    return info
            else:
                gene_loc = self.gene_loc(gene_name)
                if gene_loc:
                    if SV_res:
                        info = f'>>> \033[1;31;40m{gene_name}\033[0m : {gene_loc}\n\033[1;33;40mSV 检出\033[0m:{SV_res}\n' 
                    else:
                        info = f'>>> \033[1;31;40m{gene_name}\033[0m : {gene_loc}\n\n' 
                    return info
        if SV_res:
            info = f'\033[1;31;40m>>> 基因 {gene_name} 未找到。\n\033[0m\n\033[1;33;40mSV 检出\033[0m:{SV_res}\n' 
        else:
            info = f'\033[1;31;40m>>> 基因 {gene_name} 未找到。\n\n\033[0m'
        return info

    def CNV_search(self, CNV_file, cnv_header, Chr='', Start='', End='', gene_name=''):
        """
            CNV 查询
                Chr Start End SVSize SVType CytoBand Flag CopyNum Probes Gene GeneCount predict pli_counts
        """
        def add_res(cells, key, gene_name=''):
            value = cells[header.index(key)]
            if len(value) > 20:
                if ';' in value and gene_name:
                    genes = value.split(';')
                    value = ''
                    for i, g in enumerate(genes):
                        if g == gene_name:
                            value = f'{value},\033[1;30;43m{g}\033[0m'
                        elif i > 1:
                            break
                        else:
                            value = f'{value},{g}'
                    if f'\033[1;30;43m{gene_name}\033[0m' not in value:
                        value = f'\033[1;30;43m{gene_name}\033[0m{value} ...'
                    else:
                        value = f'{value} ...'
                else:
                    value = f'{value[0:20]} ...'
            return f'{value.strip(",")}   '
        res = ''
        with open(CNV_file, encoding='gbk') as f:
            for line in f.readlines():
                if line.startswith('Chr'):
                    header = line.strip().split('\t')
                else:
                    cells = line.strip().split('\t')
                    if gene_name and gene_name in cells[9].split(';'):
                        _, Start, End = self.gene_loc(gene_name).split('  ')
                        if int(Start) < int(cells[1]) or int(End) > int(cells[2]):
                            res = f'{res}/// [{gene_name}] \033[1;31;40m部分存在\033[0m于大片段内 : '
                        else:
                            res = f'{res}/// [{gene_name}] \033[1;31;40m整个基因存在\033[0m于大片段内 : '
                        for k in cnv_header:
                            res += add_res(cells, k, gene_name)
                        res = f'{res}\n\n'
                    elif Chr == cells[0] and Start == cells[1] and End == cells[2]:
                        SV_res = self.sv_info.call_var(f"{Chr}:{Start}-{End}")
                        for k in cnv_header:
                            res += add_res(cells, k)
                        if SV_res:
                            res = f'/// {res}\n\033[1;33;40mSV 检出\033[0m:{SV_res}\n'
                        else:
                            res = f'/// {res}\n\n'
        if not res and not gene_name:
            res = f'\033[1;31;40m/// {Chr} {Start} {End} 不在大片段内!\033[0m\n'
        return res

    @staticmethod
    def close_depth_sample(stat_file, sample_ID):
        """
            挑选深度最接近的 3 个样本, 先证者放最前面 （用于 igv 核查） 
            （英文逗号为了更方便选取所有列举的样本）
        """
        try:
            df = pd.read_table(stat_file, low_memory=False, encoding='utf-8')
        except UnicodeDecodeError:
            df = pd.read_table(stat_file, low_memory=False, encoding='gbk')
        for k in df.columns:
            if k in ['目标区域平均深度', '平均测序深度']:
                break
        df = df.sort_values(by=k, ignore_index=True).copy()
        sex = df.loc[df['样本编号'] == sample_ID, '性别'].map({'male' : '男', 'female' : '女'}).tolist()[0]
        name = ''
        if '姓名' in df.columns:
            name = df.loc[df['样本编号'] == sample_ID, '姓名'].to_list()[0]
        index = df.loc[df['样本编号'] == sample_ID,:].index[0]
        res = f'该批共 {df.shape[0]} 个样本, \033[1;31;40m深度相近\033[0m样本: '
        if df.shape[0] < 4:
            deps = "，".join([str(i) for i in df.loc[:,k].to_list()])
            samples = df.loc[:, "样本编号"].to_list()
            samples.remove(sample_ID)
            samples = [sample_ID] + samples
            res = f'{res}\033[1;30;47m{"，".join(samples)}\033[0m  {deps}\n'
        elif index == 0:
            deps = "，".join([str(i) for i in df.loc[0:2,k].to_list()])
            samples = df.loc[0:2, "样本编号"].to_list()
            samples.remove(sample_ID)
            samples = [sample_ID] + samples
            res = f'{res}\033[1;30;47m{"，".join(samples)}\033[0m  {deps}\n'
        elif index == df.shape[0]-1:
            deps = "，".join([str(i) for i in df.loc[(index-2):index,k].to_list()])
            samples = df.loc[(index-2):index, "样本编号"].to_list()
            samples.remove(sample_ID)
            samples = [sample_ID] + samples
            res = f'{res}\033[1;30;47m{"，".join(samples)}\033[0m  {deps}\n'
        else:
            deps = "，".join([str(i) for i in df.loc[(index-1):index+1,k].to_list()])
            samples = df.loc[(index-1):index+1, "样本编号"].to_list()
            samples.remove(sample_ID)
            samples = [sample_ID] + samples
            res = f'{res}\033[1;30;47m{"，".join(samples)}\033[0m  {deps}\n'
        return(sex, name, res)

    # def PO_info(self, batch, sample_ID):
    #     """ 获取样本亲属 """
    #     url = 'http://10.1.1.220:8001/adminServer?s=Bioinformatics.Api.Ngs.getAnalysisList'
    #     header = {'AUTHORIZATION-API' : '9i4t2gpa&MNFLF'}
    #     json_data = {
    #         'batch' : batch,
    #         'barcode' : [sample_ID]
    #     }
    #     res = requests.post(url , headers=header, json=json_data)
    #     if res.json()['data'] == [] or res.status_code != 200:
    #         return(f'\033[1;31;40m【错误】查询亲属样本出错\033[0m : {res.text} {json_data}\n')
    #     triosInfo = eval(res.json()['data'][0]['triosInfo'])
    #     res = sample_ID
    #     for t in triosInfo:
    #         if t['编号'] and self.sample_info(t['编号']):
    #             res = f"{t['编号']}({t['关系']})  {res}，{t['编号']}"
    #     if res == sample_ID:
    #         return f'\033[1;31;40m未找到 P0 亲属样本。\033[0m\n'
    #     else:
    #         res = res.replace(sample_ID, f'\033[1;30;47m{sample_ID}')
    #         return f'\033[1;31;40m亲属样本: \033[0m {res} \033[0m\n'

    def rel_info(self, sample_ID):
        """ 获取样本亲属 """
        url = 'http://10.1.1.220:8001/adminServer?s=Bioinformatics.Api.Ngs.getSampleRelation'
        header = {'AUTHORIZATION-API' : '9i4t2gpa&MNFLF'}
        json_data = {
            'barcode' : sample_ID
        }
        res = requests.post(url , headers=header, json=json_data)
        if res.json()['data'] == [] or res.status_code != 200:
            return(f'\033[1;31;40m【错误】查询亲属样本出错\033[0m : {res.text} {json_data}\n')
        out = sample_ID
        for info in res.json()['data']:
            if info['id'] == sample_ID:
                continue
            if not self.ngs_sample(info['id']):
                continue
            out = f"{info['id']} ({info['name']} - {info['relation']})  {out}，{info['id']}"
        if out == sample_ID:
            return f'\033[1;31;40m未找到 P0 亲属样本。\033[0m\n'
        else:
            out = out.replace(sample_ID, f'\033[1;30;47m{sample_ID}')
            return f'\033[1;31;40m亲属样本: \033[0m {out}\033[0m\n'

    @staticmethod
    def ngs_sample(sample_ID):
        """ 样本是否有二代结果 """
        url = 'http://10.10.10.150:8250/ngs/SamplePath'
        json_data = {
            'sample' : sample_ID,
            'new' : '1'
        }
        res = requests.post(url, json=json_data)
        if res.status_code != 200 or res.json()['message'] == []:
            return False
        return True


if __name__ == '__main__':
    input_str = ' '.join(sys.argv[1:])
    if input_str.strip():
        Search_var().Run(input_str)

